import numpy as np
import gradio as gr
from PIL import Image
import joblib

# 加载保存的KNN模型
model = joblib.load('best_knn_model.pkl')

# 图像预处理函数
def preprocess_image(image_data):
    # 显示原始图像
    image_data.show()

    # 将PIL图像转换为灰度模式
    pil_image = image_data.convert('L')
    
    # 打印图像信息以调试
    print("PIL image size:", pil_image.size)
    print("Image mode:", pil_image.mode)
    
    # 打印图像数据的统计信息
    image_array = np.array(pil_image)
    print("Original image array min:", image_array.min())
    print("Original image array max:", image_array.max())
    
    # 调整图像大小为8x8
    pil_image = pil_image.resize((8, 8), Image.LANCZOS)
    
    # 转换为numpy数组并调整像素值范围
    image_array = np.array(pil_image).astype(float)
    print("Resized image array min:", image_array.min())
    print("Resized image array max:", image_array.max())
    
    # 将图像数据归一化到0-16
    image_array = image_array / 255.0 * 16
    
    # 展平数组
    flattened_image = image_array.flatten()
    
    print("Preprocessed image array:", flattened_image)  # 打印预处理后的图像数据
    return flattened_image

# 预测函数
def predict_digit(image_data):
    print("Received image data:", image_data)
    
    if image_data is None:
        print("No image provided.")
        return "No image provided.", {}
    
    # 显示图像以调试
    image_data.show()
    
    # 预处理图像
    processed_image = preprocess_image(image_data)
    if processed_image is None:
        return "No valid image provided.", {}
    
    # 模型预测
    prediction = model.predict([processed_image])[0]
    probabilities = model.predict_proba([processed_image])[0]
    results = {str(i): float(prob) for i, prob in enumerate(probabilities)}
    
    return results, str(prediction)

# 创建Gradio接口
iface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Image(type="pil", label="Upload an image of a digit"),
    outputs=[gr.Label(num_top_classes=3), gr.Textbox(label="预测结果")],
    live=False,  # 关闭实时模式
    title="手写数字识别",
    description="上传一个数字图像（0-9），模型将尝试识别它。",
    allow_flagging="never",  # 禁用标记功能
)

# 启动Gradio接口
iface.launch(share=True)
